import rospy
from sensor_msgs.msg import Imu
import math
import matplotlib
from pylab import *

print "starting..."
class AccelerometerLog:
    def __init__(self):
        rospy.init_node('Acc')
        self._acc_data = [0,0,0]
        self._acc_all = 0.0
        self._Eacc = [0,0,0]
        self._time = 0.0
        self._dt = 0.0
        self._speed = 0.0
        self._ctr = 0
        self._acc_sub = rospy.Subscriber('/SENSORS/IMU',Imu, self.IMUCallback)

    def IMUCallback(self,data):
        self._dt = (data.header.stamp.to_sec() - self._time)
        self._acc_data = [data.linear_acceleration.x,data.linear_acceleration.y,data.linear_acceleration.z]
        self._acc_all = sqrt(data.linear_acceleration.x**2+data.linear_acceleration.y**2+data.linear_acceleration.z**2)
        self._time = data.header.stamp.to_sec()

        if abs(self._acc_data[0]) < 0.14 and abs(self._speed) < 0.5:
            self.Eacc()
        self.estimate_speed()

    def Eacc(self):
        for i in range(0,2):
            self._Eacc[i] = (self._Eacc[i]*self._ctr + self._acc_data[i]*1)/(self._ctr+1)
        self._ctr += 1
        if self._ctr > 120:
            self._ctr = 100

    def estimate_speed(self):
        self._speed += (self._acc_data[0] - self._Eacc[0]) * self._dt

def add2list(alist,obj,num):
    alist.append(obj)
    while len(alist) > num:
        alist.pop(0)
    return alist
    
num = 30
acc = AccelerometerLog()
print "accelerometer loger created"
Eacc_log = []
acc_log_x = []
acc_log_y = []
acc_log_z = []
acc_log_all = []
speed_log = []
figure(1)
ax1 = subplot(4,1,1)
ax2 = subplot(4,1,2)
ax3 = subplot(4,1,3)
ax4 = subplot(4,1,4)
ion()
rospy.sleep(0.25)
while(acc._ctr < 10):
    rospy.sleep(0.05)
print "collected data"
try:
    while True:
        rospy.sleep(0.05)

        add2list(Eacc_log,acc._Eacc,num)
        
        add2list(acc_log_x,acc._acc_data[0] - acc._Eacc[0],num)
        add2list(acc_log_y,acc._acc_data[1] - acc._Eacc[1],num)
        add2list(acc_log_z,acc._acc_data[2] - acc._Eacc[2],num)
                
        add2list(speed_log,acc._speed,num)

	add2list(acc_log_all,acc._acc_all,num)
            
        figure(1)
        clf()
        subplot(4,1,1)
        plot(acc_log_x)
        title("acc x")
        subplot(4,1,2)
        plot(acc_log_y)
        title("acc y")
        subplot(4,1,3)
        title('acc z')
        plot(acc_log_z)
        subplot(4,1,4)
        title('acc all')
        plot(acc_log_all)
        draw()

except:
	print "goodbye"
